from sympy import *


def Lagrange_interpolation(keys, values):
    t = len(keys)
    ploy = []
    for i in range(t):
        lst = ['((x-' + str(_) + ')/(' + str(keys[i]) + '-' + str(_) + '))' for _ in keys if _ != keys[i]]
        item = '*'.join(lst)
        ploy.append(str(values[i]) + '*' + item)
    ploy = '+'.join(ploy)

    return factor(expand(ploy))


if __name__ == '__main__':
    # example 1, interpolate a line
    x_1 = [1, 2]
    y_1 = [3, 5]
    if len(x_1) != len(y_1):
        print('The lengths of two list are not equal!')
    else:
        print('example 1 Lagrange_interpolation polynomials is:')
        print(Lagrange_interpolation(x_1, y_1))

    # example 2, interpolate a parabola
    x_2 = [0, 2, 3]
    y_2 = [1, 2, 4]
    if len(x_2) != len(y_2):
        print('The lengths of two list are not equal!')
    else:
        print('example 2 Lagrange_interpolation polynomials is:')
        print(Lagrange_interpolation(x_2, y_2))

    # example 3
    x_3 = [0, 1, 2, 3]
    y_3 = [2, 1, 0, -1]
    if len(x_3) != len(y_3):
        print('The lengths of two list are not equal!')
    else:
        print('example 3 Lagrange_interpolation polynomials is:')
        print(Lagrange_interpolation(x_3, y_3))
